#!/usr/bin/env python

# Copyright (c) 2009, Purdue University
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 
# Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or
# other materials provided with the distribution.
# 
# Neither the name of the Purdue University nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 'AS IS'
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Config export tool for Roster"""


__copyright__ = 'Copyright (C) 2009, Purdue University'
__license__ = 'BSD'
__version__ = '0.18'

import ConfigParser
import os
import sys
import shutil
import subprocess
import shlex
import socket

from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from multiprocessing import Pool
from optparse import OptionParser
from roster_config_manager import config_lib
from roster_core import config
from roster_core import errors
from roster_core import constants

# The following import allows a fake smtplib module to be used for unit tests.
if( os.getenv('ROSTERTESTPATH') is not None ):
  sys.path.insert(0,os.getenv('ROSTERTESTPATH'))
import smtplib

def EmailError(error_messages, to_address, from_address, smtp_server, subject):
  """Sends an email to report an error in exporting.

  Inputs:
    error_string: string of the body of the email
    to_address: string of email address to send error report to
  """
  error_string = '\n'.join(error_messages)
  if( None not in [error_string, to_address, from_address, smtp_server] ):
    email_message = MIMEMultipart('alternative')
    email_message['Subject'] = subject
    email_message['From'] = from_address
    email_message['To'] = to_address

    html_message = ''
    for error in error_messages:
      html_message = '%s<br/><h4>%s</h4><p>%s</p>' % (html_message, 
          error.split('\n')[0], error.lstrip('%s\n' % \
              error.split('\n')[0]).replace('\n','<br/>'))
    html_message = '<html><head></head><body>%s</body></html>' % html_message

    email_message.attach(MIMEText(error_string, 'plain'))
    email_message.attach(MIMEText(html_message, 'html'))

    smtp_handle = None
    try:
      smtp_handle = smtplib.SMTP(smtp_server)
    except socket.gaierror:
      print '%s is an invalid smtp server.' % smtp_server
    except smtplib.SMTPConnectError:
      print 'Failed to connect to %s.' % smtp_server
    if( smtp_handle is not None ):
      try:
        smtp_handle.sendmail(from_address,[to_address],
            email_message.as_string())
      except smtplib.SMTPRecipientsRefused:
        print '%s is an invalid email address.' % to_address
      smtp_handle.quit()

def RunCommandWithSingleArg(arg_dict):
  """Calls RunCommand with arguments unpacked from arg_dict.

  Inputs:
    arg_dict: dictionary of arguments to pass to RunCommand.
  Outputs:
    dictionary containing output, return code, dns_server, and command
  """
  arg_string = arg_dict['arg_string']
  dns_server = arg_dict['dns_server']
  debug = arg_dict['debug']
  command = '%s -d %s' % (arg_string, dns_server)

  return_string, return_code = RunCommand(command, debug=False)

  return {'output': return_string,
          'return_code': return_code,
          'dns_server': dns_server,
          'command': command}

def RunCommand(command, debug=False, print_command=False):
  """Runs a command and returns proper return code.

  Inputs:
    command: string of command to run
  Outputs:
    tuple of program output and return code
  """
  if( print_command ):
    print '[localhost] local: %s' % command
  process_handle = subprocess.Popen(shlex.split(str(command)), 
    stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
  stdout_output, stderr_output = process_handle.communicate()

  return_string = '%s\n%s' % (stderr_output, stdout_output)
  return_code = process_handle.returncode

  if( debug ):
    print return_string
  if( return_code is None ):
    return_code = 0
  return (return_string, return_code)

def GenerateToolArgStrings(options):
  """Generates the arg strings for dnstreeexport, dnscheckconfig,
  dnsservercheck, dnsconfigsync, and dnsquerycheck.

  Inputs:
    options: OptionParser parse's args object. Created in main()

  Outputs: 
    list of each tool's arg string.
  """
  # Preparing dnstreexport
  dnstreeexport_array = [options.tree_export]
  dnstreeexport_array.extend(['-c', options.config_file])
  if( options.force ):
    dnstreeexport_array.append('--force')
  if( options.quiet ):
    dnstreeexport_array.append('--quiet')
  dnstreeexport_arg_string = ' '.join(dnstreeexport_array)

  # Preparing dnscheckconfig
  dnscheckconfig_array = [options.check_config]
  dnscheckconfig_array.extend(['-i', '%s' % options.id])
  dnscheckconfig_array.extend(['--config-file', options.config_file])
  if( options.named_checkzone ):
    dnscheckconfig_array.extend(['-z', options.named_checkzone])
  if( options.named_checkconf ):
    dnscheckconfig_array.extend(['-c', options.named_checkconf])
  if( not options.quiet ):
    dnscheckconfig_array.append('-v')
  dnscheckconfig_arg_string = ' '.join(dnscheckconfig_array)

  # Preparing dnsservercheck
  dnsservercheck_array = [options.server_check]
  dnsservercheck_array.extend(['--export-config'])
  dnsservercheck_array.extend(['-c', options.config_file])
  dnsservercheck_array.extend(['-i', '%s' % options.id])
  dnsservercheck_arg_string = ' '.join(dnsservercheck_array)

  # Preparing dnsconfigsync
  dnsconfigsync_array = [options.config_sync]
  dnsconfigsync_array.extend(['--export-config'])
  dnsconfigsync_array.extend(['-i', '%s' % options.id])
  dnsconfigsync_array.extend(['-c', options.config_file])
  if( options.ssh_id ):
    dnsconfigsync_array.extend(['--ssh-id', options.ssh_id])
  if( options.rndc_exec ):
    dnsconfigsync_array.extend(['--rndc-exec', options.rndc_exec])
  if( options.rndc_port ):
    dnsconfigsync_array.extend(['--rndc-port', options.rndc_port])
  if( options.rndc_key ):
    dnsconfigsync_array.extend(['--rndc-key', options.rndc_key])
  if( options.rndc_conf ):
    dnsconfigsync_array.extend(['--rndc-conf', options.rndc_conf])
  dnsconfigsync_arg_string = ' '.join(dnsconfigsync_array)

  # Preparing dnsquerycheck
  dnsquerycheck_array = [options.query_check]
  dnsquerycheck_array.extend(['--export-config'])
  dnsquerycheck_array.extend(['-c', options.config_file])
  dnsquerycheck_array.extend(['-i', '%s' % options.id])
  dnsquerycheck_array.extend(['-n', '%s' % options.number])
  dnsquerycheck_array.extend(['-p', '%s' % options.port])
  dnsquerycheck_arg_string = ' '.join(dnsquerycheck_array)

  return [dnstreeexport_arg_string,
          dnscheckconfig_arg_string,
          dnsservercheck_arg_string,
          dnsconfigsync_arg_string, 
          dnsquerycheck_arg_string]

def main(args):
  """Collects command line arguments. Exports configs.

  Inputs:
    args: list of arguments from the command line
  """
  usage = ('\n'
           '\n'
           'To export database to config files:\n'
           '\t%s [-i <audit-id>] [-c <config-file>] [-q <quiet>]'
           '\n' % sys.argv[0])

  parser = OptionParser(version='%%prog (Roster %s)' % __version__, usage=usage)

  parser.add_option('-i', '--id', dest='id',
                    help='ID of tarfile output from Roster tree export.',
                    metavar='<audit-id>', default=None)
  parser.add_option('-c', '--config-file', action='store', dest='config_file',
                    help='Roster config file location.', 
                    metavar='<config-file>', 
                    default=constants.SERVER_CONFIG_FILE_LOCATION)
  parser.add_option('-q','--quiet', action='store_false', dest='quiet',
                    help='Suppress program output.', metavar='<quiet>',
                    default=False)
  parser.add_option('--tree-exporter', action='store', dest='tree_export',
                    help='Location of "dnstreeexport" binary.',
                    default='dnstreeexport')
  parser.add_option('--check-config', action='store', dest='check_config',
                    help='Location of "dnscheckconfig" binary.',
                    default='dnscheckconfig')
  parser.add_option('--server-check', action='store', dest='server_check',
                    help='Location of "dnsservercheck" binary.',
                    default='dnsservercheck')
  parser.add_option('--config-sync', action='store', dest='config_sync',
                    help='Location of "dnsconfigsync" binary.',
                    default='dnsconfigsync')
  parser.add_option('--query-check', action='store', dest='query_check',
                    help='Location of "dnsquerycheck" binary.',
                    default='dnsquerycheck')
  parser.add_option('-f', '--force', action='store_true', dest='force',
                    help='(dnstreeexport)Export trees even if nothing has '
                         'changed in the database.', default=False)
  parser.add_option('--named-checkzone', action='store',
                    dest='named_checkzone',
                    help='(dnscheckconfig)Location of named_checkzone binary.',
                    default='/usr/sbin/named-checkzone')
  parser.add_option('--named-checkconf', action='store',
                    dest='named_checkconf',
                    help='(dnscheckconfig)Location of named_checkconf binary.',
                    default='/usr/sbin/named-checkconf')
  parser.add_option('--ssh-id', action='store', dest='ssh_id',
                    help='(dnsconfigsync)SSH id file.', metavar='<ssh-id>',
                    default=None)
  parser.add_option('--rndc-exec', action='store', dest='rndc_exec',
                    help='(dnsconfigsync)Rndc executable location.',
                    metavar='<rndc-exec>', default=None)
  parser.add_option('--rndc-key', action='store', dest='rndc_key',
                    help='(dnsconfigsync)Rndc key file.', metavar='<rndc-key>',
                    default=None)
  parser.add_option('--rndc-conf', action='store', dest='rndc_conf',
                    help='(dnsconfigsync)Rndc conf file.',
                    metavar='<rndc-conf>', default=None)
  parser.add_option('--rndc-port', action='store', dest='rndc_port',
                    help='RNDC communication port.  If none provided, '
                    'named.conf will be parsed to find one.  If one can not '
                    'be found, 953 will be used.', metavar='<rndc-port>',
                    default=None)
  parser.add_option('-p','--port', action='store', dest='port',
                    help='(dnsquerycheck)Port to query DNS server on.',
                    metavar='<port>', default=53)
  parser.add_option('-n', '--number', action='store', dest='number',
                    help='(dnsquerycheck)Number of random records to query for '
                         'Default=5\nTo query all records, use \'-n all\'',
                    metavar='<number>', default=5)


  (globals()['options'], args) = parser.parse_args(args)

  if( options.config_file is None ):
    print 'The --config_file flag is required.'
    sys.exit(1)

  config_lib_instance = config_lib.ConfigLib(options.config_file)

  if( os.path.exists(config_lib_instance.root_config_dir) ):
    shutil.rmtree(config_lib_instance.root_config_dir)

  error_messages = []
  try:
    config_instance = config.Config(file_name=options.config_file)
    smtp_server = config_instance.config_file['exporter']['smtp_server']
    to_email = config_instance.config_file['exporter']['failure_'
                                                       'notification_email']
    from_email = config_instance.config_file['exporter']['system_email']
    subject = config_instance.config_file['exporter']['email_subject']
    debug = True
    if config_instance.config_file['exporter']['exporter_debug'] == 'off':
      debug = False
  except KeyError:
    print 'Incomplete config-file, error log emailing is disabled.'
    smtp_server = None
    to_email = None
    from_email = None
  
  tool_args = GenerateToolArgStrings(options)
  
  try:
    #Running dnstreeexport
    output, tree_return = RunCommand(tool_args[0], debug, print_command=True)

    if( debug and tree_return != 0 ):
      error_messages.append('%s\nReturn Code: %s\n'
          '%s' % (tool_args[0], tree_return, output))
      if( tree_return != 0 ):
        raise errors.ConfigError
   
    if( options.id is None ):
      options.id, filename = config_lib_instance.FindNewestDnsTreeFilename()

    #Since we have now set options.id, we need to regenerate.
    tool_args = GenerateToolArgStrings(options)

    #Running dnscheckconfig on all servers
    output, config_return = RunCommand(tool_args[1], debug, print_command=True)
    if( config_return != 0 ):
      raise errors.ConfigError('%s\nReturn Code: %s\n%s' % (tool_args[1], 
          config_return, output))
    config_lib_instance.UnTarDnsTree(options.id)
    all_dns_servers = config_lib_instance.FindAllDnsServers()

    #The first two tools we already ran (dnstreeexport and dnscheckconfig
    tool_args = tool_args[2:]

    for arg_string in tool_args:
      tool_name = arg_string.split(' ')[0]
      exporter_pool = Pool(processes=config_lib_instance.max_threads)
      arg_lists = []
      results = []

      for dns_server in all_dns_servers:
        arg_lists.append({'arg_string': arg_string,
                          'dns_server': dns_server,
                          'debug': debug})

      results = exporter_pool.map(RunCommandWithSingleArg, arg_lists)

      for result in results:
        output = result['output']
        return_code = result['return_code']
        dns_server = result['dns_server']
        command = result['command']
        if( debug ):
          print '[localhost] local: %s' % command

        if( return_code != 0 ):
          error_messages.append(
              '%s\nReturn Code: %s\n%s' % (
                  command, return_code, output))

          #Remove the failed dns_server. Otherwise, we'll continue to run
          #the later tools on it. 
          all_dns_servers.remove(dns_server)

  except Exception as error:
    error_messages.append(str(error))
    raise error
  finally:
    if( os.path.exists(config_lib_instance.root_config_dir) ):
      shutil.rmtree(config_lib_instance.root_config_dir)
    if( len(error_messages) > 0 ):
      EmailError(error_messages, to_email, from_email, smtp_server, subject)
      sys.exit(1)

if __name__ == '__main__':
  main(sys.argv[1:])
